﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Diagnostics;
using System.Drawing;
using System.Threading.Tasks;

using TracedReality.Logging;
using TracedReality.Primitive;
using TracedReality.RayTracing.PhotonMapping;
using TracedReality.SpatialTrees;

namespace TracedReality.RayTracing
{
    /// <summary>
    /// A ray tracer.
    /// </summary>
    public class RayTracer
    {
        public static readonly string VERSION = "1.0.0";

        private static Scene scene;
        private static PhotonMap photonmap;
        private static ShapeTree tree;
        private static List<Shape> infiniteGeometry;

        private static int totalWork;
        private static int doneWork;

        /// <summary>
        /// Gets the progress in the current render if one is active, otherwise the progress of
        /// the most recent render. May not be accurate.
        /// </summary>
        /// <returns>The progress of the render.</returns>
        public static double getProgress() { return doneWork / (double)totalWork; }

        /// <summary>
        /// Give a Scene to the RayTracer.
        /// Creates the SDS if enabled in the Settings.
        /// Creates the photon map if enabled in the Settings.
        /// </summary>
        /// <param name="s">The Scene to use for future ray tracing.</param>
        public static void giveScene(Scene s)
        {
            scene = s;

            // Make the SDS if we are supposed to use one
            if (Settings.useSDSCaster)
            {
                if (Settings.SDSType == "KD") { tree = new ShapeKDTree(scene.getBoundingBox()); }
                else if (Settings.SDSType == "Oct") { tree = new ShapeOctTree(scene.getBoundingBox()); }
                else { throw new NotImplementedException("Tree type " + Settings.SDSType); }

                // If the bounding box of a piece of geometry is infinite
                // then it should not be put in the tree and should always
                // be checked for ray intersection
                infiniteGeometry = new List<Shape>();
                foreach (Shape p in scene.shapes)
                {
                    if (p.bound.isFinite()) { tree.add(p); }
                    else { infiniteGeometry.Add(p); }
                }

                tree.split(Settings.maxSDSNodeCount, Settings.maxSDSDepth);
                tree.shrink();
                Log.log("Tree stats: { Count: " + tree.getCount() + ", Height: " + tree.getHeight() + " }");
            }
            else { tree = null; infiniteGeometry = null; }
        }

        /// <summary>
        /// Check if the given Ray intersects anything before the point along at it distance t.
        /// </summary>
        /// <param name="ray">The Ray to check.</param>
        /// <param name="t">How far along the Ray.</param>
        /// <returns>True if the Ray intersects something before t, false otherwise.</returns>
        public static bool intersectsAnythingBefore(Ray ray, double t)
        {
            if (Settings.useSDSCaster)
            {
                Hit h = new Hit();
                foreach (Shape s in infiniteGeometry)
                {
                    if (s.Intersect(ray, h) != 0)
                    {
                        if (h.t < t) { return true; }
                    }
                }

                return tree.intersectAnythingBefore(ray, t);
            }
            else
            {
                Hit h = new Hit();
                foreach (Shape p in scene.shapes)
                {
                    if (p.Intersect(ray, h) != 0)
                    {
                        if (h.t < t) { return true; }
                    }
                        
                }
            }

            return false;
        }

        /// <summary>
        /// Cast a Ray through the Scene and update the given Hit.
        /// </summary>
        /// <param name="ray">Ray to cast.</param>
        /// <param name="hit">Hit to update with intersection info.</param>
        /// <returns>1 if the Ray intersects a Shape from the outside,
        /// -1 if the Ray intersects a Shape from the inside,
        /// and 0 if no intersection occurs.</returns>
        public static int castRay(Ray ray, Hit hit)
        {
            int answer = 0;

            if (Settings.useSDSCaster)
            {
                foreach (Shape s in infiniteGeometry)
                {
                    switch (s.Intersect(ray, hit))
                    {
                        case 0:
                            break;
                        case 1:
                            answer = 1;
                            break;
                        case -1:
                            answer = -1;
                            break;
                    }
                }

                switch (tree.castRay(ray, hit))
                {
                    case 0:
                        break;
                    case 1:
                        answer = 1;
                        break;
                    case -1:
                        answer = -1;
                        break;
                }
            }
            else
            {
                foreach (Shape p in scene.shapes)
                {
                    switch (p.Intersect(ray, hit))
                    {
                        case 0:
                            break;
                        case 1:
                            answer = 1;
                            break;
                        case -1:
                            answer = -1;
                            break;
                    }
                }
            }

            return answer;
        }

        /// <summary>
        /// Recursively trace a Ray through the Scene.
        /// Terminates early if the level is too high or the contribution too low.
        /// </summary>
        /// <param name="ray">Ray to trace.</param>
        /// <param name="hit">Hit to update with intersection info.</param>
        /// <param name="level">How many bounces have occurred before this.</param>
        /// <param name="contrib">How much this Ray contributes to the result.</param>
        /// <param name="r">Random to use for sampling.</param>
        /// <returns>A Vector containing the resulting color from tracing this Ray.</returns>
        public static Vector traceRay(Ray ray, Hit hit, int level, double contrib, Random r)
        {
            hit = new Hit();
            
            if (level > Settings.numBounces) { return new Vector(); }
            if (contrib < Settings.minContribution) { return new Vector(); }

            // If it hits nothing return the background
            int result = castRay(ray, hit);
            if (result == 0 || hit.t == double.PositiveInfinity)
            {
                return Settings.backColor;
            }

            Material mat = hit.material;
            Debug.Assert(mat != null);

            hit.normal = result * hit.normal;
            Vector normal = hit.normal;
            Vector point = ray.pointAtParameter(hit.t);

            Vector answer;

            Vector diffuse;
            // Diffuse color comes from the photon map or an ambient term
            if (Settings.usePhotonMapping)
            {
                diffuse = mat.diffuse * photonmap.gatherIndirect(point, normal, ray.direction);
            }
            else
            {
                diffuse = mat.diffuse * Settings.ambientColor;
            }

            // + lighting from each light
            foreach (Light l in scene.lights)
            {
                diffuse += l.getContribution(ray, hit, r);
            }

            answer = diffuse;

            // Specular bounces
            if (mat.isSpecular())
            {
                double eta1;
                double eta2;
                if (result == 1) { eta1 = Settings.airIndex; eta2 = mat.index; }
                else { eta1 = mat.index; eta2 = Settings.airIndex; }

                double R;
                if (mat.isRefractive())
                {
                    R = Utilities.computeReflectance(ray.direction, normal, eta1, eta2);
                }
                else { R = 1; }

                double T = 1 - R;

                // Reflection
                Vector reflected;
                if (mat.isReflective() && R > 0)
                {
                    Ray reflection = new Ray(point, Vector.Reflection(ray.direction, normal));
                    Hit r_hit = null;
                    double newcontrib = mat.reflective.average() * contrib;
                    reflected = mat.reflective * traceRay(reflection, r_hit, level + 1, newcontrib, r);
                }
                else { reflected = new Vector(); }

                // Refraction
                Vector refracted;
                if (mat.isRefractive() && T > 0)
                {
                    Vector sDir = Vector.Refraction(ray.direction, normal, eta1, eta2);
                    Ray refraction = new Ray(point - (Settings.EPSILON * normal), sDir);
                    Hit s_hit = null;
                    double newcontrib2 = mat.refractive.average() * contrib;
                    refracted = mat.refractive * traceRay(refraction, s_hit, level + 1, newcontrib2, r);
                }
                else { refracted = new Vector(); }

                answer += (R * reflected) + (T * refracted);
            }

            return answer;
        }

        /// <summary>
        /// Find the color of the given pixel location through ray tracing.
        /// Performs supersampling if enabled in the Settings.
        /// </summary>
        /// <param name="x">The x coordinate of the pixel to trace.</param>
        /// <param name="y">The y coordinate of the pixel to trace.</param>
        /// <param name="maxDim">The largest dimension in the resulting image.</param>
        /// <param name="r">Random to use for sampling.</param>
        /// <returns>The color resulting from tracing the specified pixel.</returns>
        private static Vector tracePixel(int x, int y, double maxDim, Random r)
        {
            Vector pixel = new Vector();

            // Stratify for AA
            Vector[,] stratSamples = Utilities.stratify(new Vector(x, y, 0),
                new Vector(1, 0, 0),
                new Vector(0, 1, 0),
                Settings.numAntialiasSamples,
                false,
                r);

            foreach (Vector strata in stratSamples)
            {
                double xImage = (strata.x + 0.5 - (Settings.width / 2)) / maxDim + 0.5;
                double yImage = (strata.y + 0.5 - (Settings.height / 2)) / maxDim + 0.5;

                Ray ray = scene.camera.generateRay(xImage, yImage);
                Hit hit = null;
                Vector color = traceRay(ray, hit, 0, 1, r);

                // Make sure to keep this division so it is the average of all samples
                pixel += color / stratSamples.Length;
            }

            return pixel;
        }

        public static Color[,] traceArea(Utilities.Tile tile, double maxdim, Random r)
        {
            Color[,] res = new Color[tile.width, tile.height];
            int stx = tile.x;
            int sty = tile.y;

            for (int x = 0; x < tile.width; x++)
            {
                for (int y = 0; y < tile.height; y++)
                {
                    res[x, y] = (Color)tracePixel(stx + x, sty + y, maxdim, r);
                }
            }


            return res;
        }

        /// <summary>
        /// Ray trace an image of the Scene.
        /// Returns a blank image if something is missing that is required.
        /// Runs in parallel if specified in the Settings.
        /// </summary>
        /// <returns>The resulting image.</returns>
        public static Bitmap traceImage()
        {
            Bitmap image = new Bitmap(Settings.width, Settings.height, System.Drawing.Imaging.PixelFormat.Format24bppRgb);

            // Make sure everything we need exists
            if (scene == null) { Log.log("Null scene"); return image; }
            if (scene.camera == null || scene.lights == null || scene.shapes == null) { Log.log("Null scene information"); return image; }
            if (Settings.useSDSCaster && tree == null) { Log.log("No SDS"); return image; }
            //if (Settings.usePhotonMapping && photonmap == null) { Log.log("No photon map"); return image; }

            doneWork = 0;
            totalWork = 1;

            // Make the photon map if we are supposed to
            if (Settings.usePhotonMapping)
            {
                Random r = new Random();
                photonmap = new PhotonMap(scene, r);
            }
            else { photonmap = null; }

            double maxDim = Math.Max(Settings.width, Settings.height);

            if (!Settings.parallelRender)
            {
                totalWork = Settings.width;
                Random r = new Random();
                for (int x = 0; x < Settings.width; x++)
                {
                    for (int y = 0; y < Settings.height; y++)
                    {
                        Vector px = tracePixel(x, y, maxDim, r);
                        image.SetPixel(x, (Settings.height - y - 1), (Color)px);
                    }

                    doneWork = x;
                }
            }
            else
            {
                // Bitmaps are not very thread friendly so just put the results
                // in an array and copy it over later
                Color[,] data = new Color[Settings.width, Settings.height];
                //Vector[] tiles = Utilities.flatten(Utilities.tile(new Vector(0, 0, 0),
                //    new Vector(Settings.width, 0, 0),
                //    new Vector(0, Settings.height, 0),
                //    Settings.chunks));
                Utilities.Tile[] tiles = Utilities.tile(Settings.width, Settings.height, Settings.chunks);

                double xstride = Settings.width / (double)Settings.chunks;
                double ystride = Settings.height / (double)Settings.chunks;

                totalWork = tiles.Length;

                Parallel.ForEach(tiles,
                    tile =>
                    {
                        Random tr = new Random();

                        Color[,] tmp = traceArea(tile, maxDim, tr);

                        int stx = tile.x;
                        int sty = tile.y;

                        for (int x = 0; x < tile.width; x++)
                        {
                            for (int y = 0; y < tile.height; y++)
                            {
                                data[x + tile.x, y + tile.y] = tmp[x, y];
                            }
                        }

                        doneWork++;
                    });

                // Copy the results to the image
                for (int x = 0; x < Settings.width; x++)
                {
                    for (int y = 0; y < Settings.height; y++)
                    {
                        image.SetPixel(x, (Settings.height - y - 1), data[x, y]);
                    }
                }
            }

            return image;
        }
    }
}
